!##############################################################################
! RETIREMENT MAIN FILE
!
! This code is the main file to solve life-cycle problem with couples + singles,
! with portfolio choice for the RETIREMENT PERIOD 
!
! Author: Annika Bacher (annika.bacher@bi.no)        
!
! This version: June 2024
!
!##############################################################################


include "globals_ret.f90"

module retirement

        use globals_ret
        use toolbox_own
        use globals
        
        ! declare private variables in main file
        
        implicit none
        
        real*8 :: maxbndm_ret, minbndm_ret, maxbndf_ret, minbndf_ret, maxbndc_ret, minbndc_ret
        
        real*8 :: output_f_ret(3), output_m_ret(3), output_c_ret(5)
        
        integer :: inx_ret, inx2_ret(1)
        
        integer :: tr, pr, tpr, llr, zr

contains


subroutine initialize_ret

        ! ********** medical expenditures ******************

        ! medical expenditures (men and women)
        ! estimates taken from Borella et al. 2020 ("the lost ones", Table 8 in online appendix -- NOT in paper) 
        
        ! estimates refer to singles in bad health, born in the 1950s
        m_eff1f = 0.497d0 -0.101d0
        m_eff2f = -0.00634d0 +0.000672d0
        m_eff3f = 0.0000277d0

        m_eff1m = 0.497d0 -0.101d0
        m_eff2m = -0.00634d0 +0.000672d0
        m_eff3m = 0.0000277d0

        m_yf = exp(-3.8530d0-0.129d0-2.076d0+3.876d0)
        m_ym = exp(-3.8530d0-0.129d0-2.076d0+3.876d0+5.547d0)
        

        do j=1,JR
        m_efff(j)=exp((m_eff1f)*(j+retage) + m_eff2f*(j+retage)**2d0+ m_eff3f*(j+retage)**3d0)
        m_effm(j)=exp((m_eff1m-0.253d0)*(j+retage)+ (m_eff2m+0.0037d0)*(j+retage)**2d0+ (m_eff3m-0.0000177d0)*(j+retage)**3d0)
        end do
        
        ! not education dependent
        m_thetam(1)=exp(0.0d0)/5000d0      ! 5000 is normalization factor.
        m_thetam(2)=exp(0.0d0)/5000d0  
        m_thetaf(1)=exp(0.0d0)/5000d0 
        m_thetaf(2)=exp(0.0d0)/5000d0 
        
        
        ! decomposition 8
        !m_eff1f = m_eff1m
        !m_eff2f = m_eff2m
        !m_eff3f = m_eff3m
        !m_efff = m_effm
        !m_yf = m_ym
        !m_thetaf = m_thetam
        
        
        ! *********** pension payments *******************
        ! 0.6 of income in last period of working life, always referring to boom and medium transitory shock
        ! for couples (1.7 of husband's income if wife has z = 2,2 (i.e. persistent and transitory 2))
        
       do tr=1,naa
            do zr =1,nzz 
            
        if (zr < 4) then
        pen_f(zr,tr) = yfs*thetafs(tr)*efffs(JJ)*0.6d0*afs_boom(4) 
        pen_m(zr,tr) = yms*thetams(tr)*effms(JJ)*0.6d0*ams_boom(4) 
        elseif (zr > 3 .AND. zr < 7) then
        pen_f(zr,tr) = yfs*thetafs(tr)*efffs(JJ)*0.6d0*afs_boom(5) 
        pen_m(zr,tr) = yms*thetams(tr)*effms(JJ)*0.6d0*ams_boom(5) 
        elseif (zr > 6) then
        pen_f(zr,tr) = yfs*thetafs(tr)*efffs(JJ)*0.6d0*afs_boom(6) 
        pen_m(zr,tr) = yms*thetams(tr)*effms(JJ)*0.6d0*ams_boom(6) 
        end if
        
        ! comment out the correct one, depending on income state of wife
        ! if z = (1,1);(1,2);(1,3) (in coupgrid_all: if (1,1) = add +1, if (1,2) = add +2, if (1,3) = add +3)
        ! pen_c(zr,tr) = ymc*thetamc(tr)*effmc(JJ)*0.6d0*coupgrid_all(((zr-1)*nz) + (nzz*nz)*0 + 1,2)*1.7d0
        ! if z = (2,1);(2,2);(2,3) (in coupgrid_all: if (2,1) = add +1, if (2,2) = add +2, if (2,3) = add +3)
         pen_c(zr,tr) = ymc*thetamc(tr)*effmc(JJ)*0.6d0*coupgrid_all_boom(((zr-1)*nz) + (nzz*nz)*1+ 2,2)*1.7d0 ! ROBUSTNESS: INCREASE TO *2.0d0
        ! if z = (3,1),(3,2),(3,3) (in coupgrid_all: if (3,1) = add +1, if (3,2) = add +2, if (3,3) = add +3)
        ! pen_c(zr,tr) = ymc*thetamc(tr)*effmc(JJ)*0.6d0*coupgrid_all(((zr-1)*nz) + (nzz*nz)*2+ 1,2)*1.7d0
         
			end do
	    end do

         
        ! ******************** survival probability ****************************
        
         ! women       
         psif(1) = 0.986110d0
         psif(2) = 0.9848233330d0
         psif(3) = 0.9834466670d0
         psif(4) = 0.9819533330d0
         psif(5) = 0.980250d0
         psif(6) = 0.9783366670d0
         psif(7) = 0.9762966670d0
         psif(8) = 0.974140d0
         psif(9) = 0.9717866670d0
         psif(10) = 0.969070d0
         psif(11) = 0.965940d0
         psif(12) = 0.9624733330d0
         psif(13) = 0.958670d0
         psif(14) = 0.9544233330d0
         psif(15) = 0.9495366670d0
         psif(16) = 0.9439566670d0
         psif(17) = 0.9377433330d0
         psif(18) = 0.930870d0
         psif(19) = 0.9232266670d0
         psif(20) = 0.0d0
         
         
         ! men     
         psim(1) = 0.977686667d0
         psim(2) = 0.975623333d0
         psim(3) = 0.97339d0
         psim(4) = 0.97097d0
         psim(5) = 0.968243333d0
         psim(6) = 0.96524d0
         psim(7) = 0.962066667d0
         psim(8) = 0.958736667d0
         psim(9) = 0.955166667d0
         psim(10) = 0.9511533330d0
         psim(11) = 0.94663d0
         psim(12) = 0.941633333d0
         psim(13) = 0.93614d0
         psim(14) = 0.930073333d0
         psim(15) = 0.923326667d0
         psim(16) = 0.915846667d0
         psim(17) = 0.907593333d0
         psim(18) = 0.89853d0
         psim(19) = 0.888616667d0
         psim(20) = 0.0d0
         
        ! decomposition 9
         ! psif = psim
         
        ! equivalence scales during retirement
        
        ! single men 
         nu_retm = 1.0d0
              
        ! single women  
         nu_retf = 1.0d0

        ! couples
         nu_retc = 1.7d0
         
 
        write(*,*)'initialization for retirement done'
         

end subroutine

subroutine solve_vfret

    
        do l=JR,1,-1            ! age loop (has to be first loop)
        
		if (l == JR) then
		llr = l
			else
        llr = l+1
		end if        

        !$OMP PARALLEL   &
        !$OMP PRIVATE(maxbndm_ret,maxbndc_ret,maxbndf_ret,minbndc_ret,minbndf_ret,minbndm_ret,&
        !$OMP          tpr,output_m_ret,output_f_ret,output_c_ret,tr, zr, pr)
        
        !$OMP DO 


        do i=1,ncc

        
        ! start loop over remaining state variables
        
        do tr=1,naa
        
                do zr =1,nzz
                
                    do pr=1,nrisk 
                    

        ! ******** single male ********
       
       
        ! specify bounds for optimization
        minbndm_ret = cmin
        if(pr == 1) then
        maxbndm_ret = max(0d0,min(cgrid(i),cmax))
			else
        maxbndm_ret = max(0d0,min((cgrid(i)-SMFs(JJ)),cmax))
        end if 
                
                
        ! optimization
        output_m_ret = &
        golden_sing_ret(minbndm_ret,maxbndm_ret,rgrid(pr),deltas,gammam,m_ym,m_thetam(tr),m_effm(l),&
                    betam,pen_m(zr,tr),cgrid(i),VFret_m(:,zr,tr,llr), nu_retm(l),psim(l),SMFs(JJ))
                           

        cash_optret_m(pr,i,zr,tr,l) = output_m_ret(1) ! this is cprime
        Voptret_m(pr,i,zr,tr,l)     = output_m_ret(2) 
        coptret_m(pr,i,zr,tr,l)     = output_m_ret(3)
        
        
        ! ******** single female ********
       
        ! specify bounds for optimization
        minbndf_ret = cmin
        if(pr == 1) then
        maxbndf_ret = max(0d0,min(cgrid(i),cmax))
        else
        maxbndf_ret = max(0d0,min((cgrid(i)-SMFs(JJ)),cmax))
        end if 
        
        
        ! optimization
        output_f_ret = &
        golden_sing_ret(minbndf_ret,maxbndf_ret,rgrid(pr),deltas,gammaf,m_yf,m_thetaf(tr),m_efff(l),&
                    betaf,pen_f(zr,tr),cgrid(i),VFret_f(:,zr,tr,llr),nu_retf(l),psif(l),SMFs(JJ))
                            

        cash_optret_f(pr,i,zr,tr,l) = output_f_ret(1) ! this is cprime
        Voptret_f(pr,i,zr,tr,l)     = output_f_ret(2)
        coptret_f(pr,i,zr,tr,l)     = output_f_ret(3)

       ! ******** couples  ********
            
            
            do tpr=1, naa
                
       
       ! specify bounds for optimization
        minbndc_ret = cmin
        if(pr == 1) then
        maxbndc_ret = max(0d0,min(cgrid(i),cmax))              
        else
        maxbndc_ret = max(0d0,min(cgrid(i)- SMFc(JJ),cmax))
        end if
        
        
        ! optimization
        output_c_ret = &
        golden_coup_ret(minbndc_ret,maxbndc_ret,rgrid(pr),deltas,deltac,gammac,gammac,m_yf,m_ym, &
                        m_thetaf(tr),m_thetam(tpr),m_efff(l),m_effm(l),betac,betaf,betam,pen_c(zr,tpr),cgrid(i), &
                        VFret_c(:,zr,tr,tpr,llr), VFret_f(:,zr,tr,llr),VFret_m(:,zr,tpr,llr),VFret_ci(:,zr,tr,tpr,llr),&
						VFret_cp(:,zr,tr,tpr,llr),nu_retc(l),psif(l),psim(l),SMFc(JJ),pen_f(zr,tpr),pen_m(zr,tpr))
  
        ! output               
        cash_optret_c(pr,i,zr,tr,tpr,l)  = output_c_ret(1) ! this is cprime
        Voptret_c(pr,i,zr,tr,tpr,l)      = output_c_ret(2)
        coptret_c(l,zr,tr,tpr,i,pr)      = output_c_ret(3) 
        Voptret_ci(pr,i,zr,tr,tpr,l)     = output_c_ret(4)
        Voptret_cp(pr,i,zr,tr,tpr,l)     = output_c_ret(5)

                end do
                
                
        end do    ! end of nrisk loop 
        
        end do    ! end of nzz loop 

        end do    ! end of naa loop
                
        end do    ! end of ncc loop          
         
          !$OMP END DO
            !$OMP END PARALLEL                
                        

        ! Find optimal share of risky assets 
        
        ! single men
        
        do i = 1,ncc
            do zr = 1,nzz
            do tr=1,naa
            
        inx2_ret = 0
        inx_ret  = 0
    
        VFret_m(i,zr,tr,l) = maxval(Voptret_m(:,i,zr,tr,l))
        inx2_ret        = maxloc(Voptret_m(:,i,zr,tr,l))
        inx_ret         = maxval(inx2_ret,1)
        

        ! policy function for consumption
        cpolret_m(i,zr,tr,l)    = coptret_m(inx_ret,i,zr,tr,l)
          
        ! single women
        inx2_ret = 0
        inx_ret  = 0
                    
        VFret_f(i,zr,tr,l)= maxval(Voptret_f(:,i,zr,tr,l))
        inx2_ret = maxloc(Voptret_f(:,i,zr,tr,l))
        inx_ret  = maxval(inx2_ret,1)
        
        ! policy function for consumption
        cpolret_f(i,zr,tr,l)    = coptret_f(inx_ret,i,zr,tr,l)
        
            end do
            end do
        end do
        
        ! couples
        do i = 1,ncc
            do tr=1,naa
                do zr=1,nzz         
                  do tpr=1, naa

        VFret_c(i,zr,tr,tpr,l)  = maxval(Voptret_c(:,i,zr,tr,tpr,l))
        
        inx2_ret = 0
        inx_ret  = 0
        inx2_ret = maxloc(Voptret_c(:,i,zr,tr,tpr,l))
        inx_ret = maxval(inx2_ret,1)
        
        ! Define the policy function in couple
        cpolret_c(l,zr,tr,tpr,i)     = coptret_c(l,zr,tr,tpr,i,inx_ret)
        VFret_ci(i,zr,tr,tpr,l)      = Voptret_ci(inx_ret,i,zr,tr,tpr,l)
        VFret_cp(i,zr,tr,tpr,l)      = Voptret_cp(inx_ret,i,zr,tr,tpr,l)
        
                    end do
                end do
            end do
        end do
         
        
    
        write(*,'(a,i3,a)')'Age: ',l+retage,' DONE!'
        
        end do         ! end of age loop 
        
end subroutine

subroutine policy_ret

     ! ASSET RELATED POLICY FUNCTIONS


     ! single women
        
        do l=1,JR
            do tr=1,naa
               do zr=1,nzz         
                do i=1,ncc
                
        inx2_ret = 0
        inx_ret  = 0
                    
        inx2_ret = maxloc(Voptret_f(:,i,zr,tr,l))
        inx_ret  = maxval(inx2_ret,1)
        
        cprimepolret_f(i,zr,tr,l) = cash_optret_f(inx_ret,i,zr,tr,l)
        sharepolret_f(i,zr,tr,l)  = rgrid(inx_ret)
        riskypolret_f(i,zr,tr,l)  = sharepolret_f(i,zr,tr,l)*cprimepolret_f(i,zr,tr,l) 
        savepolret_f(i,zr,tr,l)   = (1d0-sharepolret_f(i,zr,tr,l))*cprimepolret_f(i,zr,tr,l)
        
                 end do
                end do
            end do
        end do


      ! single men
        
        do l=1,JR
            do tr=1,naa
               do zr=1,nzz         
                do i=1,ncc
                
        inx2_ret = 0
        inx_ret  = 0
                    
        inx2_ret = maxloc(Voptret_m(:,i,zr,tr,l))
        inx_ret  = maxval(inx2_ret,1)
                    
        cprimepolret_m(i,zr,tr,l) = cash_optret_m(inx_ret,i,zr,tr,l)
        sharepolret_m(i,zr,tr,l)  = rgrid(inx_ret)
        riskypolret_m(i,zr,tr,l)  = sharepolret_m(i,zr,tr,l)*cprimepolret_m(i,zr,tr,l) 
        savepolret_m(i,zr,tr,l)   = (1d0-sharepolret_m(i,zr,tr,l))*cprimepolret_m(i,zr,tr,l)

                 end do
                end do
            end do
        end do
        

       ! couple
        
        do l=1,JR
            do ti=1,naa
              do tpr=1,naa
                 do zr=1,nzz         
                do i=1,ncc
                
        inx2_ret = 0
        inx_ret  = 0
                      
        inx2_ret = maxloc(Voptret_c(:,i,zr,ti,tpr,l))
        inx_ret  = maxval(inx2_ret,1)
                    
        cprimepolret_c(l,zr,ti,tpr,i) = cash_optret_c(inx_ret,i,zr,ti,tpr,l)
        sharepolret_c(l,zr,ti,tpr,i)  = rgrid(inx_ret)
        riskypolret_c(l,zr,ti,tpr,i)  = sharepolret_c(l,zr,ti,tpr,i)*cprimepolret_c(l,zr,ti,tpr,i) 
        savepolret_c(l,zr,ti,tpr,i)   = (1d0-sharepolret_c(l,zr,ti,tpr,i))*cprimepolret_c(l,zr,ti,tpr,i)

                  end do
                end do
              end do
            end do
        end do
        
        
    write(*,*) 'policy functions retirement done'
    

end subroutine

subroutine hcval_ret

!single women, low education 
call comp_dist_ret(savepolret_f, riskypolret_f, pen_f(:,1), m_yf,m_efff,m_thetaf(1),1,distret_f_low)
				 
!! export results
!open(unit=1,file='results/distret_f_low.txt', &
!status='replace', action='write')

!write(1,*)distret_f_low

!do i=1,1
!close(unit=i)
!end do 

!single women, high education 
call comp_dist_ret(savepolret_f, riskypolret_f, pen_f(:,2), m_yf,m_efff, m_thetaf(2),2,distret_f_high) 

!! export results
!open(unit=1,file='results/distret_f_high.txt', &
!status='replace', action='write')

!write(1,*)distret_f_high

!do i=1,1
!close(unit=i)
!end do 

!single men, low education 
call comp_dist_ret(savepolret_m, riskypolret_m, pen_m(:,1), m_ym,m_effm, m_thetam(1),1,distret_m_low) 

!! export results
!open(unit=1,file='results/distret_m_low.txt', &
!status='replace', action='write')

!write(1,*)distret_m_low

!do i=1,1
!close(unit=i)
!end do 

!single men, high education 
call comp_dist_ret(savepolret_m, riskypolret_m, pen_m(:,2), m_ym,m_effm, m_thetam(2),2,distret_m_high) 
				
!open(unit=1,file='results/distret_m_high.txt', &
!status='replace', action='write')

!write(1,*)distret_m_high

!do i=1,1
!close(unit=i)
!end do 

end subroutine

!computing hc values
subroutine  comp_dist_ret(safepolret, riskypolret,pen,m_y, m_eff,m_theta,ind, distret)

real*8, intent(in) :: safepolret(ncc,nzz,naa,JR), riskypolret(ncc,nzz,naa,JR)

real*8, intent(in) :: pen(nzz), m_y, m_eff(JR), m_theta

integer, intent(in) :: ind

real*8, intent(out) :: distret(JR,nzz*ncc,nzz*ncc)


integer :: gridl_boom, gridl_rec, gridh_boom, gridh_rec

integer :: l, i, z, apr, rpr

real*8, dimension(JR,ncc,nzz,nrr) :: apr_boom, apr_rec
        
real*8  :: diff_boom, diff_rec 

! transition matrix across state space
distret = 0d0

do l = 1, JR-1
	do i = 1,ncc
		do z = 1,nzz
	
	do rpr = 1,nrr

!determine aprime given shock outcomes, depending on aggregrate state
apr_boom(l,i,z,rpr) = &
(1d0+r)*safepolret(i,z,ind,l) + pen(z)-(m_y*m_eff(l+1)*m_theta) + (1d0+ar_boom(rpr))*riskypolret(i,z,ind,l)

apr_rec(l,i,z,rpr) = &
(1d0+r)*safepolret(i,z,ind,l) + pen(z)-(m_y*m_eff(l+1)*m_theta) + (1d0+ar_rec(rpr))*riskypolret(i,z,ind,l)


! translate into point on asset grid
gridl_boom = max(cntarray(cgrid,ncc,apr_boom(l,i,z,rpr)),1)   
gridh_boom = min(gridl_boom+1,ncc)

gridl_rec = max(cntarray(cgrid,ncc,apr_rec(l,i,z,rpr)),1)   
gridh_rec = min(gridl_rec+1,ncc)

if (gridl_boom .NE. gridh_boom) then 
diff_boom = max((apr_boom(l,i,z,rpr)- cgrid(gridl_boom))/(cgrid(gridh_boom)-cgrid(gridl_boom)),0d0)
elseif  (gridl_boom == gridh_boom) then
diff_boom = 0d0 
end if 

if (gridl_rec .NE. gridh_rec) then 
diff_rec  =  max((apr_rec(l,i,z,rpr) - cgrid(gridl_rec)) /(cgrid(gridh_rec)-cgrid(gridl_rec)),0d0)
elseif (gridl_rec == gridh_rec) then
diff_rec = 0d0 
end if 

! always allocate to higher grid point (to avoid houshols ending up in lowest grid point with positive probability)
diff_boom = 1d0
diff_rec = 1d0

! fill matrix of probabilities
distret(l,(z-1)*ncc+i,(z-1)*ncc+gridh_boom) = distret(l,(z-1)*ncc+i,(z-1)*ncc+gridh_boom) + diff_boom*(1d0-p_rec)*pir_boom(rpr)
distret(l,(z-1)*ncc+i,(z-1)*ncc+gridl_boom) = distret(l,(z-1)*ncc+i,(z-1)*ncc+gridl_boom) + (1d0-diff_boom)*(1d0-p_rec)*pir_boom(rpr)

distret(l,(z-1)*ncc+i,(z-1)*ncc+gridh_rec) = distret(l,(z-1)*ncc+i,(z-1)*ncc+gridh_rec) + diff_rec*p_rec*pir_rec(rpr)
distret(l,(z-1)*ncc+i,(z-1)*ncc+gridl_rec) = distret(l,(z-1)*ncc+i,(z-1)*ncc+gridl_rec) + (1d0-diff_rec)*p_rec*pir_rec(rpr)

     end do
			end do
		end do
	end do


end subroutine  

subroutine VF_retirement()

          ! save the VF at the ENTRY to reitrement for
          ! single men, single women & couples
          ! this is the object to be used as cont. for working period
 

          do i=1,ncc
            do tr=1,naa
                   do zr=1,nzz         


          cont_f(i,zr,tr) = VFret_f(i,zr,tr,1)
          cont_m(i,zr,tr) = VFret_m(i,zr,tr,1)
          
          
                do tpr=1,naa
          cont_c(i,zr,tr,tpr) = VFret_c(i,zr,tr,tpr,1)
          cont_ci(i,zr,tr,tpr)= VFret_ci(i,zr,tr,tpr,1) 
          cont_cp(i,zr,tr,tpr)= VFret_ci(i,zr,tr,tpr,1) 
                end do
          
              end do
            end do
          end do
          
      write(*,*) 'specification of continuation values done'
          
end subroutine

subroutine export_ret


        ! store results in txt files (singles)
        open(unit=1,file='results/vffret.txt', &
        status='replace', action='write')
        open(unit=2,file='results/vfmret.txt', &
        status='replace', action='write')
        open(unit=3,file='results/shareret_f.txt', &
        status='replace',action='write')
        open(unit=4,file='results/shareret_m.txt', &
        status='replace',action='write')
        open(unit=5,file='results/cprimeret_f.txt', &
        status='replace',action='write')
        open(unit=6,file='results/cprimeret_m.txt', &
        status='replace',action='write')
        open(unit=7,file='results/riskyret_f.txt', &
        status='replace',action='write')
        open(unit=8,file='results/riskyret_m.txt', &
        status='replace',action='write')
        open(unit=9,file='results/saferet_f.txt', &
        status='replace',action='write')
        open(unit=10,file='results/saferet_m.txt', &
        status='replace',action='write')
        open(unit=11,file='results/consumret_f.txt', &
        status='replace',action='write')
        open(unit=12,file='results/consumret_m.txt', &
        status='replace',action='write')
        
        write(1,*)VFret_f
        write(2,*)VFret_m
        write(3,*)sharepolret_f
        write(4,*)sharepolret_m
        write(5,*)cprimepolret_f
        write(6,*)cprimepolret_m
        write(7,*)riskypolret_f
        write(8,*)riskypolret_m
        write(9,*)savepolret_f
        write(10,*)savepolret_m
        write(11,*)cpolret_f
        write(12,*)cpolret_m
        
        
        ! store results in txt files (couples)
        open(unit=13,file='results/vfcret.txt', &
        status='replace',action='write')
        open(unit=14,file='results/shareret_c.txt', &
        status='replace',action='write')
        open(unit=15,file='results/cprimeret_c.txt', &
        status='replace',action='write')
        open(unit=16,file='results/riskyret_c.txt', &
        status='replace',action='write')
        open(unit=17,file='results/saferet_c.txt', &
        status='replace',action='write')
        open(unit=18,file='results/consumret_c.txt', &
        status='replace',action='write')
       
        write(13,*)VFret_c
        write(14,*)sharepolret_c
        write(15,*)cprimepolret_c
        write(16,*)riskypolret_c
        write(17,*)savepolret_c
        write(18,*)cpolret_c
     
     
   do i=1,18
        close(unit=i)
   end do

end subroutine


end module
